import glob

import matplotlib.pyplot as plt
import numpy as np
import torch
from PIL import Image
from torch.utils.data import random_split, DataLoader

from AnimaiClassModel import MonkeyClassificationModel

img_dir = 'H:\\workspace\\pythons\\data\\'

# 加载图片并转换为数组
# Load image and convert image to multidimensional array
def image_to_array(images_folder):
    dataset = []
    for i in range(10):
        for filename in glob.glob(images_folder + "/n{}/*.jpg".format(i)):
            im = Image.open(filename)
            im = im.resize((400,300))
            pixels = np.asarray(im).astype('float32')
            pixels /= 255.0
            pixels = torch.from_numpy(pixels)
            dataset.append((pixels, i))
    return dataset


def accuracy(outputs, labels):
    _, preds = torch.max(outputs, dim=1)
    return torch.tensor(torch.sum(preds == labels).item() / len(preds))


def evaluate(model, val_loader):
    outputs = [model.validation_step(batch) for batch in val_loader]
    return model.validation_epoch_end(outputs)


def fit(epochs, lr, model, train_loader, val_loader, opt_func=torch.optim.SGD):
    history = []
    optimizer = opt_func(model.parameters(), lr)
    for epoch in range(epochs):
        # Training Phase
        for batch in train_loader:
            loss = model.training_step(batch)
            loss.backward()
            optimizer.step()
            optimizer.zero_grad()
        # Validation phase
        result = evaluate(model, val_loader)
        model.epoch_end(epoch, result)
        history.append(result)
    return history



def predict_image(input_img, model):
    inputs = input_img.unsqueeze(0)
    predictions = model(inputs)
    _, preds  = torch.max(predictions, dim=1)
    return preds[0].item()


def real_run():
    train_dataset = image_to_array(img_dir + 'training\\training')
    test_dataset = image_to_array(img_dir + 'validation\\validation')

    input_size = 300 * 400 * 3
    num_classes = 10
    # print("Input Size: ", input_size, "\nNumber of Classes: ", num_classes)
    model = MonkeyClassificationModel(input_size, num_classes)
    # res = list(model.parameters())
    # print(res)

    # View a sample Image
    img_tensor, label = train_dataset[0]
    print(img_tensor.shape)

    plt.imshow(img_tensor)
    print('Label:', label)

    # Training and Validation dataset
    val_size = round(0.2 * len(train_dataset))
    train_size = len(train_dataset) - val_size
    train_ds, val_ds = random_split(train_dataset, [train_size, val_size])

    batch_size = 16
    # Dataloaders
    train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True)
    val_loader = DataLoader(val_ds, batch_size=batch_size)
    test_loader = DataLoader(test_dataset, batch_size=batch_size)

    learning_rate = 1e-3
    history1 = fit(100, learning_rate, model, train_loader, val_loader)
    history2 = fit(100, learning_rate / 10, model, train_loader, val_loader)
    history3 = fit(100, learning_rate / 10, model, train_loader, val_loader)
    history4 = fit(100, learning_rate / 100, model, train_loader, val_loader)
    history5 = fit(100, learning_rate / 1000, model, train_loader, val_loader)

    history = history1 + history2 + history3 + history4 + history5

    accuracies = [r['val_acc'] for r in history]
    plt.plot(accuracies, '-x')
    plt.xlabel('epoch')
    plt.ylabel('accuracy')
    plt.title('Accuracy vs. No. of epochs')

    # Evaluate on test dataset
    result = evaluate(model, test_loader)
    print(result)

    # Save
    torch.save(model, img_dir + 'model\\monkey_classification.pth')

    print('over')



# real_run()